import logging
import numpy as np
import faiss
import pickle
from lazypredict.Supervised import LazyClassifier
from sklearn.model_selection import train_test_split
from util import createXY
import sys
import warnings
from tabulate import tabulate
import pandas as pd
from tqdm import tqdm

# 配置logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)

# 忽略警告
warnings.filterwarnings('ignore')


def format_results(results_df):
    """格式化结果为漂亮的表格"""
    # 重命名索引为'Model'
    results_df = results_df.reset_index()
    results_df = results_df.rename(columns={'index': 'Model'})

    # 对结果按准确率降序排序
    results_df = results_df.sort_values('Accuracy', ascending=False)

    # 只保留需要的列
    columns = ['Model', 'Accuracy', 'Balanced Accuracy', 'ROC AUC', 'F1 Score', 'Time Taken']
    results_df = results_df[columns]

    return results_df


def train_lazy_classifier():
    try:
        # 载入和预处理数据
        logging.info("开始加载和预处理数据...")
        X, y = createXY(train_folder="../data/train", dest_folder=".")

        # 转换数据类型并归一化
        X = np.array(X).astype('float32')
        faiss.normalize_L2(X)
        y = np.array(y).astype('int32')  # 确保标签为整数类型
        logging.info("数据加载和预处理完成。")

        # 数据集分割
        logging.info("开始划分数据集...")
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.25, random_state=2023
        )
        logging.info("数据集划分完成。")

        # 打印数据信息
        for name, data in [
            ("X_train", X_train), ("X_test", X_test),
            ("y_train", y_train), ("y_test", y_test)
        ]:
            logging.info(f"{name} 类型: {type(data)}, 形状: {data.shape}")

        # 初始化LazyClassifier
        logging.info("初始化LazyClassifier并开始训练模型...")
        clf = LazyClassifier(
            verbose=0,
            ignore_warnings=True,
            custom_metric=None,
            random_state=42
        )

        # 训练模型并获取结果
        results, predictions = clf.fit(X_train, X_test, y_train, y_test)

        # 格式化并显示结果
        results_formatted = format_results(results)

        # 打印进度条和结果
        print("\n模型评估结果:")
        print("-" * 80)
        # 设置浮点数格式和左对齐
        pd.set_option('display.float_format', lambda x: '{:.4f}'.format(x))
        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', None)
        print(results_formatted.to_string(index=False))

        # 获取最佳模型（使用数值比较而不是argmax）
        best_model_name = results_formatted.loc[results_formatted['F1 Score'].astype(float).idxmax(), 'Model']
        best_model = clf.models[best_model_name]

        # 保存最佳模型
        logging.info(f"\nF1分数最高的模型是: {best_model_name}")
        logging.info("保存最佳模型...")
        with open('best_model.pkl', 'wb') as f:
            pickle.dump(best_model, f)
        logging.info("模型保存完成。")

        return best_model, results_formatted

    except Exception as e:
        logging.error(f"发生错误: {str(e)}")
        raise


if __name__ == "__main__":
    try:
        best_model, results = train_lazy_classifier()
    except KeyboardInterrupt:
        logging.info("\n程序被用户中断")
    except Exception as e:
        logging.error(f"程序执行失败: {str(e)}")